# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/12/5
-------------------------------------------------
   Change Activity:
                   2022/12/5 18:01: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import json
import os

import cv2
import imageio
import numpy as np
from torch.utils.data import Dataset


def load_data(base_dir, mode, half_res=True):
    file_path = base_dir + "/transforms_" + mode + ".json"
    fp = open(file_path, "r")
    meta = json.load(fp)

    all_imgs = []
    all_poses = []

    for frame in meta["frames"]:
        img_name = os.path.join(base_dir, frame["file_path"] + ".png")
        all_imgs.append(imageio.imread(img_name))
        all_poses.append(np.array(frame["transform_matrix"]))

    all_imgs = (np.array(all_imgs) / 255).astype(np.float32)
    all_poses = np.array(all_poses).astype(np.float32)

    H, W = all_imgs[0].shape[:2]
    camera_angle_x = np.array(float(meta["camera_angle_x"]))
    focal = 0.5 * W / np.tan(0.5 * camera_angle_x)

    if half_res:
        H = H // 2
        W = W // 2
        focal = focal // 2

        all_imgs_half = np.zeros((all_imgs.shape[0], H, W, 4))
        for i, img in enumerate(all_imgs):
            all_imgs_half[i] = cv2.resize(img, (W, H), interpolation=cv2.INTER_AREA)
        all_imgs = all_imgs_half
    return all_imgs, all_poses, [H, W, focal]


class NeRFDataset(Dataset):
    def __init__(self, base_dir="./lego", mode="train"):
        super(NeRFDataset, self).__init__()
        assert mode in ["train", "test", "val"], "mode should be train ,test or val"
        self.mode = mode

        self.imgs, self.poses, self.hwf = load_data(base_dir=base_dir, mode=mode)

        self.imgs = self.imgs[..., :3] * self.imgs[..., -1:] + (1 - self.imgs[..., -1:])
        self.imgs = self.imgs[...,::-1]
        # cv2.imshow("idsa",self.imgs[0])
        # cv2.waitKey()

    def __len__(self):
        return len(self.poses)

    def __getitem__(self, index):
        return self.imgs[index], self.poses[index]


if __name__ == "__main__":
    model = NeRFDataset()
